from openai import OpenAI
from contentforge.core.model_interface import AIModelInterface, ModelConfig
from typing import List, Dict, Any, Optional

class KimiModel(AIModelInterface):
    def __init__(self, config: ModelConfig):
        super().__init__(config)
        self.client = OpenAI(api_key=config.api_key, base_url=config.api_base)

    def generate_text(
        self, 
        messages: List[Dict[str, str]], 
        system_prompt: Optional[str] = None
    ) -> str:
        try:
            # 如果有系统提示词，插入到消息列表开头
            if system_prompt:
                messages.insert(0, {"role": "system", "content": system_prompt})

            response = self.client.chat.completions.create(
                model=self.config.model_name,
                messages=messages,
                max_tokens=self.config.max_tokens,
                temperature=self.config.temperature
            )
            
            return response.choices[0].message.content.strip()
        
        except Exception as e:
            raise RuntimeError(f"Kimi模型文本生成失败: {str(e)}")

    def generate_structured_output(
        self, 
        messages: List[Dict[str, str]], 
        output_schema: Dict[str, Any],
        system_prompt: Optional[str] = None
    ) -> Dict[str, Any]:
        try:
            # 添加结构化输出的系统提示
            structured_prompt = f"""
            请按照以下JSON模式生成输出：
            {output_schema}
            
            原始消息：{messages}
            """
            
            if system_prompt:
                messages.insert(0, {"role": "system", "content": system_prompt})
            
            messages.append({"role": "user", "content": structured_prompt})

            response = self.client.chat.completions.create(
                model=self.config.model_name,
                messages=messages,
                max_tokens=self.config.max_tokens,
                response_format={"type": self.config.response_format if self.config.response_format else "text"},
                temperature=self.config.temperature
            )
            
            return response.choices[0].message.content.strip()
        
        except Exception as e:
            raise RuntimeError(f"Kimi模型结构化输出失败: {str(e)}") 